#!/bin/bash
# 可配置环境变量说明:
# DBPASSWORD - 数据库root密码 (优先使用)   export DBPASSWORD=your_password
# WAFLOGINUSER - WAF后台登录用户名 (可选)  export WAFLOGINUSER=admin  默认admin
# CONFIG_DIR - SSL证书目录 (默认/opt/1panel/secret)
# WAF_DIR - WAF安装目录 (默认/opt/docker/docker-uuwaf)
# 配置变量

CONFIG_DIR="/opt/1panel/secret"    #证书目录
WAF_DIR="/opt/docker/docker-uuwaf" #存放目录
REPO_URL="https://gitee.com/lsmir2/docker-uuwaf.git"
DEFAULT_PASSWORD="Passw0rd!"
WAF_USER="uuwaf"

# DB_USE_PWD=""
# 日志函数
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1"
}

# 检查工具是否安装
check_tools() {
    local required_tools=("docker" "git" "openssl" "sed" "python3")
    for tool in "${required_tools[@]}"; do
        if ! command -v "$tool" &>/dev/null; then
            log "错误：缺少工具 $tool，请先安装该工具。"
            return 1
        fi
    done
}

# 检查并安装Python依赖
check_python_deps() {
    log "检查Python依赖包..."

    local deps=("cryptography" "mysql-connector-python")
    local missing_deps=()

    for dep in "${deps[@]}"; do
        if ! python3 -c "import $dep" &>/dev/null; then
            missing_deps+=("$dep")
        fi
    done

    if [ ${#missing_deps[@]} -gt 0 ]; then
        log "安装缺失的Python依赖包: ${missing_deps[*]}"

        # 尝试使用pip3安装依赖
        if command -v pip3 &>/dev/null; then
            if ! pip3 install "${missing_deps[@]}"; then
                log "错误：无法安装Python依赖包，请手动安装后再运行脚本。"
                return 1
            fi
        else
            log "错误：未找到pip3，请先安装pip3后再运行脚本。"
            return 1
        fi
    else
        log "所有Python依赖包均已安装"
    fi
}

# 生成安全密码
generate_password() {
    openssl rand -base64 32 | tr -dc 'a-zA-Z0-9!@#$%^&*()_+-=' | head -c 24
}

# 克隆或更新仓库
update_repo() {
    if [ ! -d "${WAF_DIR}" ]; then
        if ! git clone "$REPO_URL"; then
            log "错误：无法克隆仓库，请检查网络连接或仓库地址是否正确。"
            return 1
        fi
    fi
}

# 初始化数据库
init_database() {
    local hostname="localhost"
    local port="3306"
    local username="root"
    local waf_db="uuwaf"
    DB_USE_PWD=$(generate_password)

    # SQL语句数组
    local sql_statements=(
        "DROP USER IF EXISTS '${WAF_USER}'@'%';"
        "DROP DATABASE IF EXISTS ${waf_db};"
        "CREATE DATABASE ${waf_db};"
        "CREATE USER '${WAF_USER}'@'%' IDENTIFIED BY '${DB_USE_PWD}';"
        "GRANT ALL PRIVILEGES ON ${waf_db}.* TO '${WAF_USER}'@'%';"
    )

    # 执行SQL
    for sql in "${sql_statements[@]}"; do
        if ! docker exec -i 1panel-mariadb mariadb -h"${hostname}" -P"${port}" -u"${username}" -p"${DBPASSWORD}" -e "${sql}"; then
            log "错误：执行SQL语句失败: ${sql}"
            return 1
        fi
    done

    # 导入SQL文件
    if ! docker cp "${WAF_DIR}/data/waf.sql" 1panel-mariadb:/docker-entrypoint-initdb.d/waf.sql ||
        ! docker exec -i 1panel-mariadb sh -c "cat /docker-entrypoint-initdb.d/waf.sql | mariadb -h${hostname} -uroot -p\"${DBPASSWORD}\" ${waf_db} && rm /docker-entrypoint-initdb.d/waf.sql"; then
        log "错误：导入waf.sql脚本失败"
        return 1
    fi

    # 环境变量WAFLOGINUSER 如果有值,则更新数据库用户登录名
    if [ -n "${WAFLOGINUSER}" ]; then
        log "更新数据库用户登录名"${WAFLOGINUSER}
        docker exec -i 1panel-mariadb mariadb -h"${hostname}" -P"${port}" -u"${username}" -p"${DBPASSWORD}" -e "UPDATE uuwaf.waf_users SET usr='${WAFLOGINUSER}' WHERE id=1;"
    fi


    echo "${WAF_USER}:${DB_USE_PWD}"
}

# 配置SSL证书
configure_ssl() {
    # 使用全局变量 WAF_USER 和 DB_USE_PWD
    log "配置SSL证书 - 用户名: ${WAF_USER}"
    log "配置SSL证书 - 密码: ${DB_USE_PWD}"


    if ! python3 ${WAF_DIR}/data/update_ssl.py -i 1 -n "${CONFIG_DIR}/server.crt" -k "${CONFIG_DIR}/server.key" -s "${WAF_USER}:${DB_USE_PWD}@127.0.0.1:3306"; then
        log "错误：updateWafssl执行失败。"
        log "python3 ${WAF_DIR}/data/update_ssl.py -i 1 -n "${CONFIG_DIR}/server.crt" -k "${CONFIG_DIR}/server.key" -s "${WAF_USER}:${DB_USE_PWD}@127.0.0.1:3306
        return 1
    fi

    if [ -f "${CONFIG_DIR}/server.crt" ] && [ -f "${CONFIG_DIR}/server.key" ]; then
        # 检查是否已经写入过配置
        if ! grep -q "${CONFIG_DIR}/server.key" "${WAF_DIR}/docker-compose.yml"; then
            echo "            - ${CONFIG_DIR}/server.key:/uuwaf/web/conf/server.key" >>"${WAF_DIR}/docker-compose.yml"
            echo "            - ${CONFIG_DIR}/server.crt:/uuwaf/web/conf/server.crt" >>"${WAF_DIR}/docker-compose.yml"
        else
            log "SSL证书配置已存在，跳过写入"
        fi
    fi
}

# 主安装函数
install_waf() {
    log "开始安装uuWAF"

    if ! check_tools; then
        exit 1
    fi

    if ! check_python_deps; then
        exit 1
    fi

    # 优先使用环境变量DBPASSWORD
    if [ -z "${DBPASSWORD}" ]; then
        read -s -p "请输入数据库root密码: " DBPASSWORD
        echo
        if [ -z "$DBPASSWORD" ]; then
            log "错误：密码不能为空"
            exit 1
        fi
    else
        log "使用环境变量DBPASSWORD作为数据库密码"
    fi

    if ! update_repo; then
        exit 1
    fi

    if ! init_database; then
        log "错误：数据库初始化失败。"
        exit 1
    fi

    if ! configure_ssl; then
        log "错误：SSL证书配置失败。"
        exit 1
    fi

    if ! cd "${WAF_DIR}"; then
        log "错误：无法进入 ${WAF_DIR} 目录。"
        exit 1
    fi

    # 使用全局变量 DB_USE_PWD
    # 转义密码中的特殊字符，避免sed命令出错
    escaped_password=$(printf '%s' "$DB_USE_PWD" | sed -e 's/[\/&]/\\&/g')
    if ! sed -i "s/uuwaf:.*@tcp/uuwaf:${escaped_password}@tcp/g" .env; then
        log "错误：无法修改.env文件。"
        exit 1
    fi

    if ! docker compose up -d; then
        if ! docker compose up -d; then
            log "错误：docker compose启动失败。"
            exit 1
        fi
    fi

    log "安装完成"
    log "数据库用户: ${WAF_USER}"
    log "数据库密码: ${DB_USE_PWD}"
    log "后台密码: ${DEFAULT_PASSWORD}"
}

# 复制SQL文件
copy_sql_file() {
    if ! docker cp 1panel-uuwaf:/uuwaf/initdb/waf.sql "${WAF_DIR}/data/waf.sql"; then
        log "错误：无法复制waf.sql文件到本地。"
        exit 1
    fi
    log "waf.sql 文件已复制到 ${WAF_DIR}/data/waf.sql"
}

# 主菜单
main() {
    # 如果有命令行参数，则直接执行对应操作
    if [ $# -gt 0 ]; then
        case "$1" in
        1) install_waf ;;
        2) copy_sql_file ;;
        *)
            log "错误：无效参数 '$1'，请使用1或2"
            exit 1
            ;;
        esac
        return
    fi

    # 否则显示交互式菜单
    echo "请选择操作:"
    echo "1. 安装uuWAF并初始化数据库"
    echo "2. 从容器复制waf.sql文件"
    read -p "请输入选项 (1或2): " choice

    case "$choice" in
    1) install_waf ;;
    2) copy_sql_file ;;
    *)
        log "错误：无效选项"
        exit 1
        ;;
    esac
}

# 传递所有参数给main函数
main "$@"
